classdef ndt_specify_training_and_test_label_remapping_DS < ndt_basic_DS
    
%  This datasource object (DS) allows one to train a classifier on a specific 
%  set of labels, and then test the classifier on a different set of 
%  labels - which enables one to evaluate how similar neural representations 
%  are across different but related conditions. This datasource
%  is a subclass of ndt_basic_DS and so it inherits all the functionality of 
%  that datasource.  
%
%  The constructor for this datasource contains contains the same arguments
%  as ndt_basic_DS, plus two additional arguments 'the_training_label_number' 
%  and 'the_test_label_numbers' i.e., the constructor has the form:  
%     ds = ndt_specify_training_and_test_label_remapping_DS(the_data, the_labels, num_cv_splits, the_training_label_numbers, the_test_label_numbers).
%  The_training_label_number and the_test_label_numbers are cell arrays that
%  specify which labels should belong to which class, with the first element
%  of these cells arrays specifying the training/test labels for the first class
%  the second element of the cell array specifies which labels belong to 
%  the second class, etc..  For example, suppose one was interested in testing
%  position invariance, and had done an experiment in which data was recorded 
%  while 10 different objects were shown at three different locations.  If the
%  labels for the 10 objects at the first location were 1:10, at the second location
%  were 11:20, and at the third location were 21:30, then one could test position
%  invariance by setting the_training_label_numbers = {1, 2, 3, 4, 5 ...}, and
%  setting the the_test_label_numbers = {[11 21], [12 22], [13, 23] ...} (it would
%  also be good to do all three variants of this (i.e., 
%  setting the_training_label_numbers = {11, 12, 13, ...} and testing with 
%  the_test_label_numbers = {[1 21], [2 22], [3, 23] ...}, etc..   The object
%  is able to test such generalization from training on one set of labels and
%  testing on a different set of labels by remapping the training label numbers to the
%  index number in the_training_label_numbers cell array, and remapping the 
%  test label numbers with the the index number into the the_test_label_numbers
%  cell array.  
%
%  There is also one additional property that can be set for this object which is:  
%  make_each_CV_split_consist_of_independent_data  (default value is 0).
%  When this argument is set to 0, the get_data method returns the normal leave one split
%  out training and test data sets (i.e., the training set consists of 
%  (num_cv_splits - 1) splits of the data and the test set consists of 1 split of the data).
%  The data in the training still comes from different splits as the data in the 
%  test set, thus one can have some of the same labels in the both 
%  the_training_label_numbers and in the_test_label_numbers (in fact, if ones has
%  the_training_label_numbers = {1, 2, 3, 4, ...}, and sets 
%  the_test_label_numbers = the_training_label_numbers, then the get_data
%  method will be the same as the ndt_basic_DS get_data method.  However, 
%  if make_each_CV_split_consist_of_independent_data = 1, then each training 
%  and test set will consist data from only split, and each cross-validation
%  split will consist of independent data.  In this case the_training_label_numbers
%  and the_test_label_numbers must not contain any of the same labels (otherwise,
%  they would be copies of the same population vector which would violate the
%  fact that the training and the test set must not have any of the same data).  
%
%  Note: label_numbers_to_use should not be set in this function (and if it is it will
%    be ignored).  Unfortunately Matlab does not allow one to overwrite set
%    method in subclasses so I can not explicitly disallow this (perhaps I 
%    should have had this class contain a ndt_basic_DS rather than be a subclass
%    of it to get around this, but keeping it as is for now).
%

%==========================================================================

%     This code is part of the Neural Decoding Toolbox.
%     Copyright (C) 2011 by Ethan Meyers (emeyers@mit.edu)
% 
%     This program is free software: you can redistribute it and/or modify
%     it under the terms of the GNU General Public License as published by
%     the Free Software Foundation, either version 3 of the License, or
%     (at your option) any later version.
% 
%     This program is distributed in the hope that it will be useful,
%     but WITHOUT ANY WARRANTY; without even the implied warranty of
%     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%     GNU General Public License for more details.
% 
%     You should have received a copy of the GNU General Public License
%     along with this program.  If not, see <http://www.gnu.org/licenses/>.
    
%==========================================================================  



properties 
    
    the_training_label_numbers = [];   % a cell array specifying which label numbers should be used for training the classifier
                                       %   i.e., the_training_label_numbers = {[class_1_training_nums], [class_2_training_nums], ...}
    the_test_label_numbers = [];       % a cell array specifying which label numbers should be used for testing the classifier
                                       %   i.e., the_training_label_numbers = {[class_1_test_nums], [class_2_test_nums], ...}
    
    make_each_CV_split_consist_of_independent_data = 0;  % if this is set to 1 then each CV splits consists has independent data, i.e.,
                                                         %  each CV split has the amount of data that is typically only in the test set, 
                                                         %  and every CV training set does not consist of (num_cv_splits -1) * num_labels data points
                                                         %  but instead consists of length(cell2mat(the_test_label_numbers)) training points.
                                       
end


    

methods 

    
    function ds = ndt_specify_training_and_test_label_remapping_DS(the_data, the_labels, num_cv_splits, the_training_label_numbers, the_test_label_numbers)
    
       ds = ds@ndt_basic_DS(the_data, the_labels, num_cv_splits);   % set properties using parent constructor
                
       ds.the_training_label_numbers = the_training_label_numbers;   % set properties that are unique to this object
       ds.the_test_label_numbers = the_test_label_numbers;
          
       % get all the unique labels
       if length(the_training_label_numbers) ~= length(the_test_label_numbers)
            error('The cell arrays the_training_label_numbers and the_test_label_numbers must have the same number of cells (with the vector in cell i containing which labels belong to class i)');
       end
       

    end
    

       
    function the_properties = get_DS_properties(ds)    
    
        the_properties = get_DS_properties@ndt_basic_DS(ds);
        
        the_properties.the_training_label_numbers = ds.the_training_label_numbers;
        the_properties.the_test_label_numbers  = ds.the_test_label_numbers; 
        the_properties.make_each_CV_split_consist_of_independent_data = ds.make_each_CV_split_consist_of_independent_data;
        
    end
    
    
    
    function  [XTr_all_time_cv YTr_all_cv XTe_all_time_cv YTe_all_cv ADDITIONAL_DATASOURCE_INFO] = get_data(ds)
      
        
        the_training_label_numbers = ds.the_training_label_numbers;
        the_test_label_numbers = ds.the_test_label_numbers;
                
        % only use labels that are listed in the_training_label_numbers and the_test_label_numbers
        % if label_numbers_to_use has been set it will be ignored...
        the_training_nums = cell2mat(the_training_label_numbers);
        the_test_nums = cell2mat(the_test_label_numbers);
        label_numbers_to_use = [the_training_nums(:); the_test_nums(:)];
        %label_numbers_to_use = union(cell2mat(the_training_label_numbers), cell2mat(the_test_label_numbers)); 
        ds.label_numbers_to_use = label_numbers_to_use;
        
        
        [XTr_all_time_cv YTr_all_cv XTe_all_time_cv YTe_all_cv ADDITIONAL_DATASOURCE_INFO] = get_data@ndt_basic_DS(ds);
      
        
        
        if ds.make_each_CV_split_consist_of_independent_data == 1
            
            
            % if running each training and test set separately, there can not be any overlap between the labels listed in the training and test sets 
            %(otherwise there could be some of the same data in the training and test sets which is completely forbidden!!!!!
            all_unique_training_labels = unique(cell2mat(the_training_label_numbers));
            all_unique_test_labels = unique(cell2mat(the_test_label_numbers));

    
            if length(intersect(all_unique_training_labels,  all_unique_test_labels)) > 0
                error('if running each split separately (i.e., make_each_CV_split_consist_of_independent_data == 1), then none of the same labels can be in the training and test sets, otherwise there will be some of the same data in the training and test sets')
            end
        
            
            % remap labels (only using old 'test' data because I want each CV split to be independent)
            for iCV = 1:length(XTr_all_time_cv{1})

                    remapped_YTr_all_cv{iCV} = NaN .* ones(size(YTe_all_cv{iCV}));  % only getting labels (and data) from test set b/c want each CV split to be independent
                    remapped_YTe_all_cv{iCV} = NaN .* ones(size(YTe_all_cv{iCV}));

                    for iGroup = 1:length(the_training_label_numbers)
                        remapped_YTr_all_cv{iCV}(ismember(YTe_all_cv{iCV}, the_training_label_numbers{iGroup})) = iGroup;
                        remapped_YTe_all_cv{iCV}(ismember(YTe_all_cv{iCV}, the_test_label_numbers{iGroup})) = iGroup;
                    end
            end
            
            
            
            % remove data from trials in which the labels are not appropriate for the training/test sets
            for iCV = 1:length(XTr_all_time_cv{1})

               train_inds  = ~isnan(remapped_YTr_all_cv{iCV});
               test_inds  = ~isnan(remapped_YTe_all_cv{iCV});
          

               for iTime = 1:length(XTr_all_time_cv) 
                   XTr_all_time_cv{iTime}{iCV} = XTe_all_time_cv{iTime}{iCV}(:, train_inds);   % only getting data from test set b/c want each CV split to be independent
                   XTe_all_time_cv{iTime}{iCV} = XTe_all_time_cv{iTime}{iCV}(:, test_inds);
               end
     
               YTr_all_cv{iCV} = remapped_YTr_all_cv{iCV}(train_inds);   % remove NaNs from remapped labels
               YTe_all_cv{iCV} = remapped_YTe_all_cv{iCV}(test_inds);
               
            end
   
            

        else
        
            % remap labels
            for iCV = 1:length(XTr_all_time_cv{1})

                    remapped_YTr_all_cv{iCV} = NaN .* ones(size(YTr_all_cv{iCV}));
                    remapped_YTe_all_cv{iCV} = NaN .* ones(size(YTe_all_cv{iCV}));

                    for iGroup = 1:length(the_training_label_numbers)
                        remapped_YTr_all_cv{iCV}(ismember(YTr_all_cv{iCV}, the_training_label_numbers{iGroup})) = iGroup;
                        remapped_YTe_all_cv{iCV}(ismember(YTe_all_cv{iCV}, the_test_label_numbers{iGroup})) = iGroup;
                    end
            end

            % remove data from trials in which the labels are not appropriate for the training/test sets
            for iCV = 1:length(XTr_all_time_cv{1})

               training_data_inds_to_remove = isnan(remapped_YTr_all_cv{iCV});
               test_data_inds_to_remove = isnan(remapped_YTe_all_cv{iCV});


               for iTime = 1:length(XTr_all_time_cv) 
                   XTr_all_time_cv{iTime}{iCV}(:, training_data_inds_to_remove) = [];
                   XTe_all_time_cv{iTime}{iCV}(:, test_data_inds_to_remove) = [];
               end

               remapped_YTr_all_cv{iCV}(training_data_inds_to_remove) = [];
               remapped_YTe_all_cv{iCV}(test_data_inds_to_remove) = [];

            end

            YTr_all_cv = remapped_YTr_all_cv;
            YTe_all_cv = remapped_YTe_all_cv;
        
        end
        
        
        %size(XTr_all_time_cv{1}{1})
        %size(XTe_all_time_cv{1}{1})

        
        
    end  % end get_data method
        
    
    
    % It would be good if it impossible to set label_numbers_to_use from an outside call 
    %  but I can't do this b/c it is not possible to overwrite set methods in subclasses.
    %  So now if anyone sets label_numbers_to_use this function call will be ignored (which is not great).
    %  I could work around this by having this class contain an instance of basic_DS instead of inheriting from it (something to consider).
    
    
end   % end methods




end % end class
